#include "addresses.h"
#include "common.h"
#include "conversion.h"
#include "kerl.h"
#include <string.h>
#include "keccak/sha3.h"
#define CHECKSUM_CHARS 9

//extern Conversion* pConversion;
extern Keccak384* pKeccak384;

namespace Address {
	namespace {
		void digestSingleChunk(unsigned char *key_fragment,
				Keccak384 *digest_sha3, Keccak384 *round_sha3)
		{
			for (int k = 0; k < 26; k++) {
				Kerl::initialize(round_sha3);
				Kerl::absorbChunk(round_sha3, key_fragment);
				Kerl::squeezeFinalChunk(round_sha3, key_fragment);
			}

			// absorb buffer directly to avoid storing the digest fragment
			Kerl::absorbChunk(digest_sha3, key_fragment);
		}

	// initialize the sha3 instance for generating private key
		void initSHAs(const unsigned char *seed_bytes, uint32_t idx,
							  Keccak384 *key_sha, Keccak384 *digest_sha,
							  unsigned char *buffer)
		{
			// use temp bigint so seed not destroyed
			os_memcpy(buffer, seed_bytes, NUM_HASH_BYTES);

			Conversion conv;
			conv.bytes_add_u32_mem(buffer, idx);

			Kerl::initialize(key_sha);
			Kerl::absorbChunk(key_sha, buffer);
			Kerl::squeezeFinalChunk(key_sha, buffer);

			Kerl::initialize(key_sha);
			Kerl::absorbChunk(key_sha, buffer);

			Kerl::initialize(digest_sha);
		}
	}

	void getAddressTrytes(const unsigned char *seed_bytes, uint32_t idx, unsigned int security, char *address) {
		unsigned char bytes[48];
		getPublicAddr(seed_bytes, idx, security, bytes);
		Conversion conv;
		conv.bytes_to_chars(bytes, address, 48);
	}

	// generate public address in byte format
	void getPublicAddr(const unsigned char *seed_bytes, uint32_t idx,
						 unsigned int security, unsigned char *address_bytes)
	{
		if (!IN_RANGE(security, MIN_SECURITY_LEVEL, MAX_SECURITY_LEVEL)) {
			THROW(INVALID_PARAMETER);
		}

		Keccak384 digest_sha;	// !important ... don't use FPGA for nested SHA3s

		// buffer for the digests of each security level
		unsigned char digest[NUM_HASH_BYTES * security];

		// use last chunk of digest, as this is only used after the key is generated
		unsigned char *buffer = digest + NUM_HASH_BYTES * (security - 1);

		// init private key sha, digest sha
		initSHAs(seed_bytes, idx, pKeccak384, &digest_sha, buffer);

		for (uint8_t i = 0; i < security; i++) {
			for (uint8_t j = 0; j < 27; j++) {
				// use address output array as a temp Kerl state storage
				unsigned char *state = address_bytes;

				// the state takes only 48bytes and allows us to reuse key_sha
				Kerl::stateSqueezeChunk(pKeccak384, state, buffer);
				// re-use key_sha as round_sha
				digestSingleChunk(buffer, &digest_sha, pKeccak384);

				// as key_sha has been tainted, reinitialize with the saved state
				Kerl::reinitialize(pKeccak384, state);
			}
			Kerl::squeezeFinalChunk(&digest_sha, digest + NUM_HASH_BYTES * i);

			// reset digest sha for next digest
			Kerl::initialize(&digest_sha);
		}

		// absorb the digest for each security
		Kerl::absorbBytes(&digest_sha, digest, NUM_HASH_BYTES * security);

		// one final squeeze for address
		Kerl::squeezeFinalChunk(&digest_sha, address_bytes);
	}

	// get 9 character checksum of NUM_HASH_TRYTES character address
	void getAddressWithChecksum(const unsigned char *address_bytes,
								   char *full_address)
	{
		Kerl::initialize(pKeccak384);

		unsigned char checksum_bytes[NUM_HASH_BYTES];
		Kerl::absorbChunk(pKeccak384, address_bytes);
		Kerl::squeezeFinalChunk(pKeccak384, checksum_bytes);

		char full_checksum[NUM_HASH_TRYTES];
		Conversion conv;
		conv.bytes_to_chars(checksum_bytes, full_checksum, NUM_HASH_BYTES);

		conv.bytes_to_chars(address_bytes, full_address, NUM_HASH_BYTES);

		os_memcpy(full_address + NUM_HASH_TRYTES,
				  full_checksum + NUM_HASH_TRYTES - NUM_CHECKSUM_TRYTES,
				  NUM_CHECKSUM_TRYTES);
	}
}
